import Crypto
from Crypto.Cipher import AES
import hashlib  # a python 2.3 and 2.4 version can be found at http://code.krypto.org/python/hashlib/
import random
import struct
import exceptions

class streamcrypt:

    def __init__(self, password, version=0):
        # generate a 256 bit key from the arbitrary length password.
        # this will allow us to use AES256
        h = hashlib.sha256()
        h.update(password)
        self.keymatter = h.digest()

        self.version = version
        self.footer_length = 9

    def decStart(self):
        self.decbuf = ''
        self.ptxtbuf = ''
        self.bytes_decrypted = 0
        self.header_bytes = 0
        self.header_incomplete = 1
        self.decrypting = 1
        self.deccipher = None # will be instantiated in decNext once we have read the header

    def decNext(self, chunk):
        if self.decrypting == 0:
            raise WrongUsage('called decNext() without first calling decStart()')
        
        self.decbuf += chunk
        
        if self.header_incomplete:
            self.__read_header()
            
        if self.header_incomplete:
            return ''

        if self.deccipher == None:
            self.deccipher = AES.new(self.keymatter, AES.MODE_CBC, self.header['IV'])

        if len(self.decbuf)<16:
            return ''
        
        chunk_size = len(self.decbuf) - (len(self.decbuf) % 16) # a chunk has to be a multiple of 16 bytes in length
        self.ptxtbuf += self.deccipher.decrypt(self.decbuf[0:chunk_size])
        self.decbuf = self.decbuf[chunk_size:]

        # keep a suitable number of plaintext bytes in case they are the last ones that contain
        # the length field
        # the amount of bytes ot keep is self.footer_length + 15 because that's enough to contain the footer plus
        # any padding data that may follow
        keeplength = self.footer_length + 15
        if len(self.ptxtbuf) <= keeplength:
            return ''

        rtxt = self.ptxtbuf[:len(self.ptxtbuf)-keeplength]
        self.ptxtbuf = self.ptxtbuf[len(self.ptxtbuf)-keeplength:]

        self.bytes_decrypted += len(rtxt)
        
        return rtxt
        
    def decEnd(self):
        if self.header_incomplete:
            raise DecryptionFailed("Could not find a header in the encrypted file.")
        
        # remove the padding data from the back of the string
        # the first character of padding data should be a #
        # the remaining ones should not be
        while(len(self.ptxtbuf)>0 and self.ptxtbuf[-1] != '#'):
            self.ptxtbuf = self.ptxtbuf[:-1]

        # get the padding marker
        if(len(self.ptxtbuf) == 0):
            raise DecryptionFailed("Could not find padding marker at the end of the plaintext.")

        self.ptxtbuf = self.ptxtbuf[:-1]
        
        # get the number of bytes in the file
        if(len(self.ptxtbuf) < 8):
            raise DecryptionFailed("Could not find footer in the plaintext.")

        bytes_expected = self.__decode_length(self.ptxtbuf[-8:])
        self.ptxtbuf = self.ptxtbuf[:-8]

        self.bytes_decrypted += len(self.ptxtbuf)
        
        if(self.bytes_decrypted != bytes_expected):
            raise DecryptionFailed("Saw %i plaintext bytes but footer says there should have been %i." % (self.bytes_decrypted, bytes_expected))

        self.decrypting = 0
            
        return self.ptxtbuf
        
    def encStart(self):
        self.encbytes = 0L
        self.returned_header = 0
        self.encbuf = ''
        self.encrypting = 1

        # make a random initial value. This will be stored
        # in the header of the encrypted file.
        self.IV = ''
        for i in range(AES.block_size):
            self.IV += struct.pack('B', random.randint(0, 0xff)) 

        # create the cipher object itself
        self.enccipher = AES.new(self.keymatter, AES.MODE_CBC, self.IV)

    def encNext(self, raw):
        if self.encrypting == 0:
            raise WrongUsage('called encNext() without first calling encStart()')
        ctxt = ''
        if not self.returned_header:
            ctxt = self.__header()
            self.returned_header = 1
            
        self.encbuf += raw
        self.encbytes += len(raw)
        if len(self.encbuf)>16:
            chunk_size = len(self.encbuf) - (len(self.encbuf) % 16) # a chunk has to be a multiple of 16 bytes in length
            ctxt += self.enccipher.encrypt(self.encbuf[0:chunk_size])
            self.encbuf = self.encbuf[chunk_size:]

        return ctxt

    def encEnd(self):
        # handle header for zero length files
        ctxt = ''
        if not self.returned_header:
            ctxt += self.__header()
            self.returned_header = 1

        # append the footer
        # append a length field to the end of the clear text
        self.encbuf += self.__length_field()
        # also append a character that allows finding the end of the length field
        self.encbuf += '#'
        
        # pad data to a multiple of 16 bytes
        while len(self.encbuf) % 16 != 0:
            self.encbuf += '0'
        ctxt += self.enccipher.encrypt(self.encbuf)

        self.encrypting = 0
        return ctxt

    def __header(self):
        return struct.pack('B', self.version)+self.IV

    def __read_header(self):
        # byte 0 contains the version number
        if self.header_bytes == 0 and len(self.decbuf) > 17:
            self.header = {}
            self.header['version'] = struct.unpack('B', self.decbuf[0])[0]
            self.decbuf = self.decbuf[1:]
            self.header_bytes += 1

            if(self.header['version'] != 0):
                raise UnknownVersion(self.header['version'])

            # get the initial value for the AES algorithm (16 bytes)
            self.header['IV'] = self.decbuf[:16]
            self.decbuf = self.decbuf[16:]
            self.header_incomplete = 0

    def __length_field(self):
        # split the long integer into eight 8 bit values. This way we are absoultely sure to
        # get the same 64 bit representation no matter what the platform
        bit = 0
        bits = 64
        str = ''
        while bit < bits:
            len_byte = (self.encbytes & long(0xFF)<<bit) >> bit  # extract bits bit - bit+8
            bit += 8
            str += struct.pack('B', len_byte)

        return str

    def __decode_length(self, raw):
        # read the long 64 bit integer that represents the number of encrypted plaintext
        # bytes to be expected
        bit = 0
        bytes = 8
        length = 0
        byte = 0
        while byte < bytes:
            bit = byte*8
            length += struct.unpack('B', raw[byte])[0] << bit # read bits bit - bit+8
            byte += 1

        return length


class WrongUsage(exceptions.Exception):
    def __init__(self, reason):
        self.reason = reason
        return

    def __str__(self):
        print " ",self.reason

class DecryptionFailed(exceptions.Exception):
    def __init__(self, reason):
        self.reason = reason
        return

    def __str__(self):
        print " ","Encryption failed because "+self.reason


class UnknownVersion(exceptions.Exception):
    def __init__(self, version):
        self.version = version
        return

    def __str__(self):
        print "","Don't understand files of version %s!" % self.version

